import tkinter as tk
from tkinter import filedialog, messagebox, ttk
import numpy as np
import torch
from PIL import Image, ImageTk
from scipy.fft import fft2, ifft2, fftshift, ifftshift
from skimage import io
import inference as pre


class App(ttk.Frame):
    def __init__(self, parent):
        ttk.Frame.__init__(self, parent)
        self.root = parent
        self.root.title("手势识别")

        self.original_image = None
        self.modified_image = None
        self.image_path = None

        for index in [0, 1, 2]:
            self.columnconfigure(index=index, weight=1)
            self.rowconfigure(index=index, weight=1)

        self.current_str = "dark"

        # 创建窗口
        self.setup_widgets()

    def setup_widgets(self):
        menu_bar = tk.Menu(self.root)
        file_menu = tk.Menu(menu_bar, tearoff=0)
        file_menu.add_command(label="打开", command=self.open_image)
        file_menu.add_separator()
        file_menu.add_command(label="退出", command=self.quit)
        menu_bar.add_cascade(label="文件", menu=file_menu)
        self.root.config(menu=menu_bar)

        self.check_frame = ttk.LabelFrame(self, text="command", padding=(20, 10))
        self.check_frame.grid(
            row=0, column=0, padx=(20,10), pady=(30, 10), sticky="nsew"
        )

        self.button = ttk.Button(self.check_frame, text="predict", command=self.output_text)
        self.button.grid(row=0, column=0, padx=5, pady=10, sticky="nsew")

        # 切换主题
        self.switch = ttk.Checkbutton(
            self.check_frame, text="Switch", style="Switch.TCheckbutton", command=self.switch_theme
        )
        self.switch.grid(row=1, column=0, padx=5, pady=20, sticky="nsew")

        self.button_1 = ttk.Button(self.check_frame, text="quit", command=self.quit)
        self.button_1.grid(row=2, column=0, padx=5, pady=10, sticky="nsew")

        # 分隔窗格 #1
        self.paned = ttk.PanedWindow(self)
        self.paned.grid(row=0, column=2, pady=(10, 5), sticky="nsew", rowspan=3)

        # 窗口 #1
        self.pane_1 = ttk.Frame(self.paned, padding=5)
        self.paned.add(self.pane_1, weight=1)

        self.notebook_0 = ttk.Notebook(self.pane_1)
        self.notebook_0.pack(fill="both", expand=True)

        # 原始图像控件
        self.tab_1 = ttk.Frame(self.notebook_0)
        self.notebook_0.add(self.tab_1, text="img")
        # 创建画布
        self.canvas = tk.Canvas(self.tab_1)
        self.canvas.pack(side="left", fill="both", expand=True)

        # 创建分隔窗格 #2
        self.pane_2 = ttk.Frame(self.paned, padding=5)
        self.paned.add(self.pane_2, weight=3)

        # 创建选项卡框架
        self.notebook = ttk.Notebook(self.pane_2)
        self.notebook.pack(fill="both", expand=True)

        self.tab_1 = ttk.Frame(self.notebook)
        for index in [0, 1]:
            self.tab_1.columnconfigure(index=index, weight=1)
            self.tab_1.rowconfigure(index=index, weight=1)
        self.notebook.add(self.tab_1, text="predicted")

        self.text_box = tk.Text(self.tab_1, wrap="word", width=20, height=3)
        self.text_box.grid(row=0, column=0, padx=10, pady=10, sticky="nsew")
        self.text_box.config(state=tk.DISABLED)  # 禁用文本框编辑功能

    # 打开图像
    def open_image(self):
        file_path = filedialog.askopenfilename(filetypes=[("Image files", "*.jpg;*.jpeg;*.png;*.bmp")])
        self.image_path = file_path
        if file_path:
            try:
                self.original_image = io.imread(file_path)
                self.display_image(self.original_image, self.canvas)  # 显示原始图像
            except Exception as e:
                self.output_text(f"Failed to load image: {e}")

    # 将图像打印
    def display_image(self, image, target_canvas):
        # 调整图像尺寸以适应画布
        canvas_width = target_canvas.winfo_width()
        canvas_height = target_canvas.winfo_height()
        image_resized = Image.fromarray(image).resize((canvas_width, canvas_height))

        # 将PIL图像转换为ImageTk对象
        image_tk = ImageTk.PhotoImage(image_resized)

        # 在画布上显示图像
        target_canvas.create_image(0, 0, anchor="nw", image=image_tk)
        target_canvas.image = image_tk  # 保持对图像的引用

    # 切换主题
    def switch_theme(self):
        # 保存当前窗口大小
        current_size = self.master.winfo_geometry()

        # 更改主题
        if self.current_str == "dark":
            self.current_str = "light"
        else:
            self.current_str = "dark"

        # 调用Tcl代码来设置主题
        self.master.tk.call("set_theme", self.current_str)

        # 恢复之前的窗口大小
        self.master.geometry(current_size)

    def output_text(self):
        # 如果图像路径为空，则返回
        # 防止窗口初始化时崩溃
        if self.image_path == None:
            return
        # 清空文本框内容
        self.text_box.config(state=tk.NORMAL)
        self.text_box.delete('1.0', tk.END)
        self.text_box.config(state=tk.DISABLED)
        # 加载训练好的模型
        model = torch.load('./models/model.pkl')
        if torch.cuda.is_available():
            device = torch.device("cuda")
        else:
            device = torch.device("cpu")
        model.to(device)
        # 显示图片，输出预测结果
        predicted = pre.inference(self.image_path, model, device)
        self.text_box.config(state=tk.NORMAL)  # 先启用文本框以插入文本
        self.text_box.insert(tk.END, str(predicted) + "\n")
        self.text_box.see(tk.END)
        self.text_box.config(state=tk.DISABLED)  # 再禁用文本框以防止用户编辑

    # 退出
    def quit(self):
        if messagebox.askokcancel("退出", "确定要退出吗?"):
            self.root.destroy()

